import os
import numpy as np
import torch
import utils.utils as utils


class deterministic_agent:
    def __init__(
        self,
        device,
        env,
        Predictor,
    ):
        self.device = device
        self.env = env
        self.Predictor = Predictor
        self.model_dir = "./checkpoints"
        os.makedirs(self.model_dir, exist_ok=True)

        self.name = "deterministic_agent"

    def take_action(self, state):
        state_blc = self.process_state_to_blc(
            state,
            self.env.get_window(),
            self.env.get_future_count(),
            self.env.get_tech_num(),
        ).to(self.device)
        state_blc = utils.process_blc_features(state_blc)
        logits = self.Predictor(state_blc).squeeze(-1)
        probs = torch.sigmoid(logits)  # (0, 1)
        action = probs * 2 - 1  # (-1, 1)
        action = torch.clamp(action, -1, 1)  # 保证范围内
        return action

    def process_state_to_blc(self, state, window, future_count, tech_num):
        """
        将 state 转换为 (future_num, window, 1 + tech_num) 的格式
        确保正确处理 tech_ary 的 [len × (future_num*tech_num)] 存储结构

        Args:
            state: 输入状态 [available_amount, future_positions, price_window, tech_window]
            window: 时间窗口长度
            future_count: 期货品种数量
            tech_num: 每个期货的技术指标数量

        Returns:
            torch.Tensor: (future_num, window, 1 + tech_num)
        """
        expected_len = (
            1 + future_count + window * future_count + window * future_count * tech_num
        )
        if len(state) != expected_len:
            raise ValueError(f"State长度错误，预期{expected_len}，实际{len(state)}")
        ptr = 0
        ptr += 1
        future_positions = state[ptr : ptr + future_count]  # future_count个值
        ptr += future_count
        price_size = window * future_count
        price_flat = state[ptr : ptr + price_size]
        ptr += price_size
        price = price_flat.reshape(window, future_count).T
        tech_flat = state[ptr:]
        tech = tech_flat.reshape(window, future_count, tech_num).transpose(1, 0, 2)
        price = price[:, :, np.newaxis]  # (future_count, window, 1)
        state_blc = np.concatenate(
            [price, tech], axis=-1
        )  # (future_count, window, 1+tech_num)
        return self.norm_blc(torch.tensor(state_blc, dtype=torch.float32))

    def norm_blc(self, x):
        means = x.mean(1, keepdim=True).detach()
        x = x - means
        stdev = torch.sqrt(torch.var(x, dim=1, keepdim=True, unbiased=False) + 1e-5)
        x /= stdev
        return x

    def eval(self):
        self.Predictor.eval()
